from stock.utils.get_stock_price import get_stock_price

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import os


# 2. LSTM预测模型
class StockPricePredictor(nn.Module):
    def __init__(self, input_size=6, hidden_size=16):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True,
            dropout=0.2
        )
        self.fc = nn.Sequential(
            nn.Linear(hidden_size, 8),
            nn.ReLU(),
            nn.Linear(8, 1)  # 直接预测标准化后的价格
        )

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        return self.fc(lstm_out[:, -1, :])  # 取最后一个时间步


# 4. 预测器
class StockPricePredictorWrapper:
    def __init__(self, model_path='saved_models'):
        self.device = torch.device('cpu')
        self.model, self.scaler_params = self._load_model(model_path)
        self.lookback = self.scaler_params['lookback']

    def _load_model(self, model_path):
        """加载模型和scaler参数"""
        checkpoint = torch.load(
            os.path.join(model_path, 'best_model.pth'),
            map_location=self.device
        )

        # 初始化模型结构
        config = checkpoint['model_config']
        model = StockPricePredictor(
            input_size=config['input_size'],
            hidden_size=config['hidden_size']
        )
        model.load_state_dict(checkpoint['model_state'])
        model.to(self.device)
        model.eval()

        return model, checkpoint['scaler_params']

    def denormalize(self, normalized_price):
        """将标准化价格反归一化"""
        return normalized_price * self.scaler_params['scale_'][1] + self.scaler_params['mean_'][1]

    def predict(self, stock_a_values, stock_b_values):
        """
        预测次日收盘价
        :param stock_a_values: 股票A最近lookback天的收盘价列表
        :param stock_b_values: 股票B最近lookback天的收盘价列表
        :return: 预测的次日B股收盘价
        """
        if len(stock_a_values) != self.lookback or len(stock_b_values) != self.lookback:
            raise ValueError(f"需要输入{self.lookback}天的历史数据")

        # 归一化输入
        raw_input = np.hstack([
            np.array(stock_a_values).reshape(-1, 1),
            np.array(stock_b_values).reshape(-1, 1)
        ])
        scaled_input = (raw_input - self.scaler_params['mean_']) / self.scaler_params['scale_']
        input_seq = scaled_input.flatten()

        # 预测
        with torch.no_grad():
            input_tensor = torch.FloatTensor(input_seq).unsqueeze(0).unsqueeze(1).to(self.device)
            predicted_norm = self.model(input_tensor).item()

        # 反归一化
        return self.denormalize(predicted_norm)


# 示例使用
if __name__ == "__main__":
    # 模拟数据
    code = '101.HG00Y'
    list1 = get_stock_price(code, 101, 3, '20250508')
    print(list1)
    code = '1.600362'
    list2 = get_stock_price(code, 101, 3, '20250508')
    print(list2)

    stock_a = [x.close for x in list1]
    stock_b = [x.close for x in list2]

    predictor = StockPricePredictorWrapper()

    predicted_price = predictor.predict(stock_a, stock_b)
    print(f"- 预测: {predicted_price:.2f}")
